#!/usr/bin/env python

import shutil
import click
import os
import python.helpers as helpers
import python.version
from python.helpers import print_v
from python.run import load_my_inequalities, run_with_inequalities


@click.group()
@click.option("--steps", default=100)
@click.option("--step-size", default=1)
@click.option("--block-size", default=None)
@click.option("--plot/--no-plot", default=False)
@click.option("--save-path", default="bs_code/data")
@click.option("--save-dir", default=None)
@click.option("--save-data/--no-save-data", default=True)
@click.option("--pack-data/--no-pack-data", default=False)
@click.option("--delete-uncompressed/--no-delete-uncompressed", default=False)
@click.option("--archive-type", default="gztar")
@click.option("--plot-type", default="pdf")
@click.option("--plot-max", default=None, type=int)
@click.option("--histogram-bins", default=50, type=int)
@click.option("--run-reduction/--estimate-only", default=False)
@click.option("--max-beta", default=30, type=int)
@click.option("--additional-fplll-params", default="", type=str)
@click.option("--use-best-step/--use-last-step", default=False)
@click.option("--max-enumerate", default=1, type=int)
@click.option("--verbose/--silent", default=True)
@click.option("--thresholds", nargs=21, type=float,
              default=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                       0.0, 0.0])  # --------
@click.option("--sweep-efforts", nargs=21, type=int,
              default=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])  # -------
@click.pass_context
def main(
        ctx,
        steps,
        step_size,
        block_size,
        plot,
        save_path,
        save_dir,
        save_data,
        pack_data,
        delete_uncompressed,
        archive_type,
        plot_type,
        plot_max,
        histogram_bins,
        run_reduction,
        max_beta,
        additional_fplll_params,
        use_best_step,
        max_enumerate,
        verbose,
        thresholds,
        sweep_efforts,
):
    ctx.ensure_object(dict)
    ctx.obj["steps"] = steps
    ctx.obj["step_size"] = step_size
    ctx.obj["block_size"] = block_size
    ctx.obj["plot"] = plot
    ctx.obj["save_data"] = save_data
    ctx.obj["pack_data"] = pack_data
    ctx.obj["delete_uncompressed"] = delete_uncompressed
    ctx.obj["archive_type"] = archive_type
    ctx.obj["plot_type"] = plot_type
    ctx.obj["histogram_bins"] = histogram_bins
    ctx.obj["plot_max"] = plot_max
    ctx.obj["run_reduction"] = run_reduction
    ctx.obj["max_beta"] = max_beta
    ctx.obj["additional_fplll_params"] = additional_fplll_params
    ctx.obj["use_best_step"] = use_best_step
    ctx.obj["max_enumerate"] = max_enumerate
    ctx.obj["save_path"] = save_path
    ctx.obj["save_dir"] = save_dir
    ctx.obj["thresholds"] = thresholds
    ctx.obj["sweep_efforts"] = sweep_efforts
    helpers.set_verbose(verbose)


@main.command()
@click.pass_context
@click.argument("path")
@click.argument("no_ineqs", type=int)
@click.argument("bilateral", type=bool)
@click.argument("version", type=int)
@click.option("--thread-count", type=int, default=12)
@click.option("--narrowing-path", default=None)
@click.option("--group", type=int, default=0)
def myload(ctx, path, no_ineqs, bilateral, version, thread_count, narrowing_path=None, group=0):
    steps = ctx.obj["steps"]
    step_size = ctx.obj["step_size"]
    plot = ctx.obj["plot"]
    save_data = ctx.obj["save_data"]
    pack_data = ctx.obj["pack_data"]
    delete_uncompressed = ctx.obj["delete_uncompressed"]
    archive_type = ctx.obj["archive_type"]
    plot_type = ctx.obj["plot_type"]
    histogram_bins = ctx.obj["histogram_bins"]
    plot_max = ctx.obj["plot_max"]
    python.version.set_version(version)
    if bilateral:
        with open('output.txt', 'w') as file:
            file.write("Process as bilateral inequalities")
            file.write("\n")
            print_v("Process as bilateral inequalities")
    else:
        with open('output.txt', 'w') as file:
            file.write("Process as unilateral inequalities")
            file.write("\n")
            print_v("Process as unilateral inequalities")
    propagation_data = load_my_inequalities(path, no_ineqs, bilateral, narrowing_path, group)
    success, propagation_data = run_with_inequalities(
        propagation_data,
        steps,
        step_size,
        ctx.obj["block_size"],
        run_reduction=ctx.obj["run_reduction"],
        max_beta=ctx.obj["max_beta"],
        add_fplll=ctx.obj["additional_fplll_params"],
        use_best_step=ctx.obj["use_best_step"],
        max_enum=ctx.obj["max_enumerate"],
        thread_count=thread_count,
        thresholds=ctx.obj["thresholds"],
        sweep_efforts=ctx.obj["sweep_efforts"],
    )
    propagation_data.set_settings(ctx.obj)
    with open('output.txt', 'a') as file:
        file.write(f"Success: {success}\n")
        print(f"Success: {success}")
    process_propagation_data(
        propagation_data,
        plot,
        save_data,
        pack_data,
        delete_uncompressed,
        plot_type,
        plot_max,
        histogram_bins,
        archive_type,
        ctx.obj["save_path"],
        ctx.obj["save_dir"],
    )


def process_propagation_data(
        data,
        plot,
        save_data,
        pack_data,
        delete_uncompressed,
        plot_type,
        plot_max,
        histogram_bins,
        archive_type,
        save_path,
        save_dir,
):
    if save_data:
        if save_path:
            data.set_dir_prefix(save_path)
        if save_dir:
            data.set_dir(save_dir)
        data.save_data()
    if plot:
        data.plot(plot_type=plot_type, bins=histogram_bins, plot_max=plot_max)
    if (plot or save_data) and pack_data:
        print_v("Packing archive..")
        os.makedirs("archives", exist_ok=True)
        shutil.make_archive(
            f"archives/{data.dir.replace('/', '_')}", archive_type, data.dir
        )
        if delete_uncompressed:
            print_v("Removing uncompressed data..")
            shutil.rmtree(data.dir)


if __name__ == "__main__":
    main(obj={})
